/**
 * @file soft_plus_v2_custom.cpp
 *
 * Copyright (C) 2023-2025. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */

#include "register/op_def_registry.h"
#include "softmax_v2_tiling.h"

const uint32_t perElementByBlk = 8; 
uint32_t getExponent(uint32_t dimR) {
  uint32_t exponent = 0;
  uint32_t tmpDimR = dimR;
  
  while (1) {
      // 计算 perElementByBlk 的 (exponent+1) 次方
      uint32_t power = 1;
      for (int32_t i = 0; i < exponent + 1; i++) {
          power *= perElementByBlk;
      }
      
      if (power > tmpDimR) {
          return exponent;
      }
      exponent++;
  }
  return exponent;
}

uint32_t getScaleFactor(uint32_t dimR, uint32_t exponent) {
  // 计算 perElementByBlk 的 exponent 次方
  uint32_t power = 1;
  for (int32_t i = 0; i < exponent; i++) {
      power *= perElementByBlk;
  }
  
  return dimR / power;
}

uint32_t Align(uint32_t size, uint32_t alignFactor) {
  if (size % alignFactor != 0) {
      return (size / alignFactor + 1) * alignFactor;
  }
  return size;
}

namespace optiling {
const uint32_t maxRSizeFloat = 2048;
const uint32_t maxRSizeHalf = 4096;

static ge::graphStatus TilingFunc(gert::TilingContext *context) {

  SoftmaxV2TilingData tiling;

  tiling.set_perElementByBlk(perElementByBlk);

  auto src_origin_shape = context->GetInputShape(0)->GetOriginShape();
  // auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
  uint64_t ub_size = 196608;
  // ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);

  uint32_t oriASize = src_origin_shape[0];
  uint32_t oriRSize = src_origin_shape[1];

  auto coreNum =oriASize >= 16 ? 16: oriASize;

  //对A轴进行核间切分
  uint32_t multiCoreASize = oriASize / coreNum;
  uint32_t tailMultiCoreASize = multiCoreASize;
  uint32_t multiCoreNum = coreNum;
  uint32_t tailMultiCoreNum = 0;
  if(oriASize % coreNum != 0){
    multiCoreNum = oriASize % coreNum;
    tailMultiCoreNum = coreNum - multiCoreNum;
    multiCoreASize += 1;
  }

  tiling.set_oriRSize(oriRSize);
  tiling.set_oriASize(oriASize);

  tiling.set_multiCoreNum(multiCoreNum);
  tiling.set_multiCoreASize(multiCoreASize);
  tiling.set_tailMultiCoreASize(tailMultiCoreASize);


  //对R轴进行核内切分
  //切分的原理在于，将2048作为R轴的主轴，其余轴作为尾轴
  //为了方便尾轴的计算，需要进一步将尾轴划分为主轴和尾轴
  uint32_t rSize = 0;
  uint32_t rloopLimit = 0;
  if(oriRSize >= maxRSizeFloat){
    rSize = maxRSizeFloat;
    rloopLimit = oriRSize / rSize;
  }
  uint32_t tailRSize = oriRSize - (rloopLimit * rSize);
  
  //计算rSize的幂(exponent)和缩放因子(scaleFactor)
  uint32_t exponent = 0;
  uint32_t scaleFactor = 0;
  if(rSize > 0){
    exponent = getExponent(rSize);
    if(exponent > 0 ){
      scaleFactor = getScaleFactor(rSize, exponent);
    }

  }

  uint32_t tailExponent = 0;
  uint32_t tailScaleFactor = 0;
  if(tailRSize > 0){
    tailExponent = getExponent(tailRSize);
    if(tailExponent > 0 ){
      tailScaleFactor = getScaleFactor(tailRSize, tailExponent);
    }
  }

  uint32_t mainR = tailScaleFactor * (1LL << (3 * tailExponent));
  uint32_t tailR = tailRSize - mainR;
  uint32_t tailRAlign8 = Align(tailR, perElementByBlk);

  tiling.set_rSize(rSize);
  tiling.set_rloopLimit(rloopLimit);
  tiling.set_tailRSize(tailRSize);
  tiling.set_exponent(exponent);
  tiling.set_scaleFactor(scaleFactor);
  tiling.set_tailExponent(tailExponent);
  tiling.set_tailScaleFactor(tailScaleFactor);
  tiling.set_mainR(mainR);
  tiling.set_tailR(tailR);
  tiling.set_tailRAlign8(tailRAlign8);

  //为了确定aloopLimit，需要确定使用的Buffer空间大小
  //srcTBuf + reduceTempTBuf + softmaxMaxTBuf * 2 + softmaxExpMaxTBuf;
  uint32_t count = 0;
  uint32_t aSize = multiCoreASize;
  uint32_t reduceTempTBufSize = 0;
  uint32_t srcTBufSize = 0;
  uint32_t softmaxMaxTBufSize = 0;
  uint32_t softmaxExpMaxTBufSize = 0;
  uint64_t totalTBufSize = 0;
  while(1){
    aSize = aSize >> count;
    //计算reduceTempTBuf
    reduceTempTBufSize = Align(aSize, perElementByBlk);
    //在有mainR的前提下，为了保全原始数据，需要先将mainR里的数据搬运到tempBuffer中，因此在mainR存在时，tempBuffer的空间至少要能存储mainR
    if(mainR > 0){
      reduceTempTBufSize = aSize * mainR;
    }
    if(rSize > 0 && mainR < scaleFactor * (1LL << (3 * (exponent - 1)))){
      reduceTempTBufSize = aSize * scaleFactor * (1LL << (3 * (exponent - 1)));
    }
  
    srcTBufSize = aSize * (mainR + tailRAlign8);
    if(rSize > 0){
      srcTBufSize = aSize * rSize;
    }
    softmaxMaxTBufSize = Align(multiCoreASize, perElementByBlk) * perElementByBlk * 2;
    softmaxExpMaxTBufSize = Align(aSize, perElementByBlk) * perElementByBlk;
  
    totalTBufSize = (reduceTempTBufSize + srcTBufSize + softmaxMaxTBufSize + softmaxExpMaxTBufSize) * sizeof(float);
    if(totalTBufSize > ub_size){
      count += 1;
    }else{
      break;
    }
  }

  uint32_t multiCoreAloopLimit = multiCoreASize / aSize;
  uint32_t multiCoreTailA = multiCoreASize - (multiCoreAloopLimit * aSize);
  tiling.set_multiCoreA(aSize);
  tiling.set_multiCoreTailA(multiCoreTailA);
  if(multiCoreTailA > 0){
    tiling.set_multiCoreAloopLimit(multiCoreAloopLimit + 1);
  }else{
    tiling.set_multiCoreAloopLimit(multiCoreAloopLimit);
  }

  tiling.set_tailMultiCoreA(0);
  tiling.set_tailMultiCoreTailAloopLimit(0);
  tiling.set_tailMultiCoreTailA(0);

  if(tailMultiCoreASize > 0){
    aSize = tailMultiCoreASize;
    count = 0;
    while(1){
      aSize = aSize >> count;
      //计算reduceTempTBuf
      reduceTempTBufSize = Align(aSize, perElementByBlk);
      //在有mainR的前提下，为了保全原始数据，需要先将mainR里的数据搬运到tempBuffer中，因此在mainR存在时，tempBuffer的空间至少要能存储mainR
      if(mainR > 0){
        reduceTempTBufSize = aSize * mainR;
      }
      if(rSize > 0 && mainR < scaleFactor * (1LL << (3 * (exponent - 1)))){
        reduceTempTBufSize = aSize * scaleFactor * (1LL << (3 * (exponent - 1)));
      }
    
      srcTBufSize = aSize * (mainR + tailRAlign8);
      if(rSize > 0){
        srcTBufSize = aSize * rSize;
      }
      softmaxMaxTBufSize = Align(multiCoreASize, perElementByBlk) * perElementByBlk * 2;
      softmaxExpMaxTBufSize = Align(aSize, perElementByBlk) * perElementByBlk;
    
      totalTBufSize = (reduceTempTBufSize + srcTBufSize + softmaxMaxTBufSize + softmaxExpMaxTBufSize) * sizeof(float);
      if(totalTBufSize > ub_size){
        count += 1;
      }else{
        break;
      }
    }
    uint32_t tailMultiCoreTailAloopLimit = tailMultiCoreASize / aSize;
    uint32_t tailMultiCoreTailA = tailMultiCoreASize - (tailMultiCoreTailAloopLimit * aSize);
    tiling.set_tailMultiCoreA(aSize);
    tiling.set_tailMultiCoreTailA(tailMultiCoreTailA);

    if(tailMultiCoreTailA > 0){
      tiling.set_tailMultiCoreTailAloopLimit(tailMultiCoreTailAloopLimit + 1);
    }else{
      tiling.set_tailMultiCoreTailAloopLimit(tailMultiCoreTailAloopLimit);
    }

  }

  reduceTempTBufSize = 1;
  //在有mainR的前提下，为了保全原始数据，需要先将mainR里的数据搬运到tempBuffer中，因此在mainR存在时，tempBuffer的空间至少要能存储mainR
  if(mainR > 0){
    reduceTempTBufSize = mainR;
  }
  if(rSize > 0 && mainR < scaleFactor * (1LL << (3 * (exponent - 1)))){
    reduceTempTBufSize = scaleFactor * (1LL << (3 * (exponent - 1)));
  }

  srcTBufSize = mainR + tailRAlign8;
  if(rSize > 0){
    srcTBufSize = rSize;
  }

  tiling.set_reduceTempTBufSize(reduceTempTBufSize);
  tiling.set_srcTBufSize(srcTBufSize);

  context->SetBlockDim(coreNum);

  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

  size_t *currentWorkspace = context->GetWorkspaceSizes(1);
  currentWorkspace[0] = 0;

  return ge::GRAPH_SUCCESS;
}
} // namespace optiling

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext *context) {
  const gert::Shape *src_shape = context->GetInputShape(0);
  gert::Shape *dst_shape = context->GetOutputShape(0);
  *dst_shape = *src_shape;
  return ge::GRAPH_SUCCESS;
}

static ge::graphStatus InferDataType(gert::InferDataTypeContext *context) {
  const auto inputDataType = context->GetInputDataType(0);
  context->SetOutputDataType(0, inputDataType);
  return ge::GRAPH_SUCCESS;
}
} // namespace ge

namespace ops {
class SoftmaxV2 : public OpDef {
public:
  explicit SoftmaxV2(const char *name) : OpDef(name) {
    this->Input("src")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT})
        .Format({ge::FORMAT_ND});

    this->Output("dst")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT})
        .Format({ge::FORMAT_ND});

    this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);

    this->AICore()
        .SetTiling(optiling::TilingFunc)
        .AddConfig("ascend910")
        .AddConfig("ascend910b")
        .AddConfig("ascend310p")
        .AddConfig("ascend310b");
  }
};
OP_ADD(SoftmaxV2);
} // namespace ops
